# https://gitee.com/yueyinqiu5990/tj12413601/blob/master/assignment3/question1/main.py
import matplotlib.pyplot
import numpy
import scipy.constants

import mediums


def _calculate_and_plot(
        subplot: int,
        title: str,
        n: tuple[float, float, float],
        h: tuple[float, float, float]):
    a = mediums.Medium.create(h[0], n[0])
    b = mediums.Medium.create(h[1], n[1])
    c = mediums.Medium.create(h[2], n[2])
    medium = mediums.Medium.join([a, b] * 4 + [a, c] + [a, b] * 5)

    matplotlib.pyplot.subplot(subplot)

    omega_0 = scipy.constants.c / (n[0] * h[0] + n[1] * h[1])
    omegas = numpy.linspace(0.5 * omega_0, 1.5 * omega_0, 100)
    ts = []
    for omega in omegas:
        transfer_matrix = medium.transfer_matrix(omega, 0)
        t = 2 / numpy.sum(transfer_matrix)
        ts.append(numpy.abs(t) ** 2)

    matplotlib.pyplot.title(title)
    matplotlib.pyplot.xlabel(r"$\omega$")
    matplotlib.pyplot.ylabel(r"$|t|^2$")
    matplotlib.pyplot.plot(omegas, ts)


def main():
    _calculate_and_plot(
        231, r"$h_c = 4275.6 \times 10^{-9}$",
        (2.6, 2.45, 1.8), (740e-9, 1329e-9, 4275.6e-9))
    _calculate_and_plot(
        232, r"$h_c = 3206.7 \times 10^{-9}$",
        (2.6, 2.45, 1.8), (740e-9, 1329e-9, 3206.7e-9))
    _calculate_and_plot(
        233, r"$h_c = 2137.8 \times 10^{-9}$",
        (2.6, 2.45, 1.8), (740e-9, 1329e-9, 2137.8e-9))
    _calculate_and_plot(
        234, r"$h_c = 1068.9 \times 10^{-9}$",
        (2.6, 2.45, 1.8), (740e-9, 1329e-9, 1068.9e-9))
    _calculate_and_plot(
        235, r"$h_c = 534.45 \times 10^{-9}$",
        (2.6, 2.45, 1.8), (740e-9, 1329e-9, 534.45e-9))
    _calculate_and_plot(
        236, r"$h_c = 356.3 \times 10^{-9}$",
        (2.6, 2.45, 1.8), (740e-9, 1329e-9, 356.3e-9))

    matplotlib.pyplot.subplots_adjust(
        left=0.079, bottom=0.083, right=0.957, top=0.926,
        wspace=0.233, hspace=0.333)
    matplotlib.pyplot.show()


if __name__ == "__main__":
    main()
